import numpy as np
from pyE17 import utils as U
from pyE17 import io
import ptycho
import os

import matplotlib.pyplot as plt
import matplotlib.cm as cm
import scipy

# Main parameters
energy = 6.2
z = 7.046        # Distance chosen to have exactly 32 nm pixel size
ds = 172e-6      # what is this
lam = 1.2398e-9 / energy

noise_levels = [1e5, 3e5, 1e6, 3e6, 1e7, 3e7, 1e8, 3e8]

# Gold, thick:
#ior = -1j*1.24638818E-05 - 7.9982543E-05 
#max_thickness = 1e-6

# Carbon, thin:
#ior = -1j*3.4054267E-08 - 1.19337546E-05
#max_thickness = 200e-9

# Carbon, thicker
ior = -1j * 3.4054267E-08 - 1.19337546E-05
max_thickness = 1e-6

run_name = 'medweak'

probe_file = 'probe_0.h5'
object_file = 'object_0.h5'

probe = io.h5read(probe_file)['probe']
probe /= U.norm(probe)
ob0 = io.h5read(object_file)['object']
object = np.exp(-2j * np.pi * ior * max_thickness * (255 - ob0) / 255. / lam)

asize = probe.shape
dx_spec = lam * z / (asize[0] * ds) 

nth = 5
dr = 0.7e-6
lx = 12e-6
ly = 12e-6

positions = ptycho.round_scan_ROI_positions(dr, lx, ly, nth)

# For now: round positions to nearest integer
positions = np.array(positions) / dx_spec
pos_min = positions.min(axis=0)
positions -= pos_min
positions = np.round(positions)
Npos = len(positions)

# This is the way the final object size is computed 
pos_max = positions.max(axis=0)
object_size = tuple(np.array(asize) + pos_max)

# Setup object views for computation of diffraction patterns
offset = (300, 300)
obj = object[offset[0]:(offset[0] + object_size[0]), offset[1]:(offset[1] + object_size[1])]
obview = [obj[p[0]:(p[0] + asize[0]), p[1]:(p[1] + asize[1])] for p in positions]

#for v, i in zip(obview, range(1, len(obview))):
#    scipy.misc.imsave('outfile%d.jpg' % i, v.real)

#obview = [object[(offset[0] + p[0]):(offset[0] + p[0] + asize[0]), (offset[1] + p[1]):(offset[1] + p[1] + asize[1])] for p in positions]
coverage = np.zeros_like(obj).astype(float)
for p in positions:
    coverage[p[0]:(p[0] + asize[0]), p[1]:(p[1] + asize[1])] += abs(probe) ** 2

io.h5write('simul_%s_solution.h5' % run_name, obj=obj.copy(), probe=probe, coverage=coverage)

Iprobe = (np.abs(np.fft.fftn(probe)) ** 2).sum()

# Loop through noise levels
for k, noise in enumerate(noise_levels):
    data = np.zeros((Npos,) + asize, dtype=int)
    for iob in range(len(obview)):
        dp = np.fft.fftshift(np.fft.fftn(probe * obview[iob]))
        data[iob, :, :] = np.random.poisson(np.abs(dp) ** 2 * noise / Iprobe)
    #sstr = ('%1.0e' % noise).replace('e+0', 'e')
    dirname = 'simul_%s_%02d' % (run_name, k)
    try:
        os.mkdir(dirname)
    except OSError:
        pass
    filename = dirname + ('/simul_%s_%02d_data_%dx%d.h5' % (run_name, k, asize[0], asize[1]))
    print filename
    io.h5write(filename, data=data, fmask=np.ones(asize, dtype=int))
